/*
aso.cpp: Antiviral Sequence Optimizer implementation.
    Copyright (C) 2009  Daniel Gutson, FuDePAN

    This file is part of Aso.

    Aso is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    Aso is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with Aso.  If not, see <http://www.gnu.org/licenses/>.

    NOTE: This file is in prototype stage, and is under active development.
*/

#define MAX_ONLY

#include <string>
//#include <vector>
#include <map>
#include <list>
#include <set>
#include <iostream>

#include <mili/mili.h>
#include <fxp.h>

#include <stl-debug/stl_debug.h>
#include <stl-debug/stl_debug_impl.h>

using namespace std;

class SparseSequence
{
STL_DEBUG_PRIVATE:
    map<size_t, char> sequence;
public:
    bool amino_at(size_t position, char& amino) const
    {
        const map<size_t, char>::const_iterator it = sequence.find(position);
        const bool ret = (it != sequence.end());

        if (ret)
            amino = it->second;

        return ret;
    }

    // this assumes that, if the position is not specified, the amino isn't what's
    // being asked. This is because resistant positions are not the original sequence,
    // and this only works for sparse sequences applicable to this particular problem.
    bool amino_is(size_t position, char amino) const
    {
        const map<size_t, char>::const_iterator it = sequence.find(position);

        if (it == sequence.end())
            return false;
        else
            return it->second == amino;
    }

    SparseSequence& set_amino(size_t position, char amino)
    {
        sequence[position] = amino;
        return *this;
    }

    bool operator < (const SparseSequence& other) const
    {
        return sequence < other.sequence;
    }

    bool operator == (const SparseSequence& other) const
    {
        return sequence == other.sequence;
    }
};

typedef list<SparseSequence> SparseSequencesList;

struct SparseAlignment: set<SparseSequence>
{
    SparseAlignment()
    {
        // there always has to be at least 1 sequence.
        insert(SparseSequence());
    }

    SparseAlignment(const SparseAlignment& other)
        : set<SparseSequence>(other)
    {}
};

class Antiviral
{
    string name;

STL_DEBUG_PRIVATE:
    struct ResistancePosition
    {
        const size_t position;
        const string aminoacids;

        ResistancePosition(size_t position, const string& aminoacids)
            : position(position),
              aminoacids(aminoacids)
        {}

        bool applicable(const SparseSequence& seq) const
        {
            char amino;
            // if the sequence doesn't have an aminoacid at that position,
            // we assume it's not the one we're looking for.
            if (seq.amino_at(position, amino))
                return (aminoacids.find(amino) == string::npos);
            else
                return true;
        }
    };

    list<ResistancePosition> resistances;
private:

    bool apply(const SparseSequence& sequence, SparseSequencesList& resistant_sequences) const;

public:
    Antiviral(const string& name)
        : name(name)
    {}

    Antiviral(const Antiviral& other)
        : name(other.name)
    {}


    Antiviral& add_resistance(size_t pos, const string& aminoacids)
    {
        resistances.push_back(ResistancePosition(pos, aminoacids));
        return *this;
    }

    const string& get_name() const
    {
        return name;
    }

    bool applicable_to(const SparseSequence& sequence) const;
    bool apply(const SparseAlignment& alig_original, SparseAlignment& alig_resulting) const;
};

typedef list<Antiviral> Antivirals;
typedef list<const Antiviral*> Therapy;

#include "aso_debug.h"

void print_therapy(const SparseAlignment& final_alignment, const Antivirals& antivirals, const Therapy& final_therapy)
{
    static size_t last_max = 0;

#ifdef MAX_ONLY
    if (final_therapy.size() >= last_max)
#endif
    {
        cout << "New final therapy: length " << (last_max = final_therapy.size()) << endl;
        cout << "\t";
        for (Therapy::const_iterator it = final_therapy.begin(); it != final_therapy.end(); ++it)
        {
            cout << (*it)->get_name() << " ";
        }
        cout << endl;
    }
}

bool is_final_sequence(const SparseSequence& seq, const Antivirals& antivirals)
{
    bool is_final = true;
    Antivirals::const_iterator it = antivirals.begin();

    while (it != antivirals.end() && is_final)
    {
        is_final = !it->applicable_to(seq);
        ++it;
    }

    return is_final;
}

bool alignment_contains_final_sequence(const SparseAlignment& alignment, const Antivirals& antivirals)
{
    bool is_final = false;
    SparseAlignment::const_iterator it = alignment.begin();
    while (it != alignment.end() && !is_final)
    {
        is_final = is_final_sequence(*it, antivirals);
        ++it;
    }

    return is_final;
}

void generate_therapy(const SparseAlignment& initial_alignment, const Antivirals& antivirals, Therapy& current_therapy)
{
    SparseAlignment new_alignment;

    if (alignment_contains_final_sequence(initial_alignment, antivirals))
    {
        // death :(
        print_therapy(initial_alignment, antivirals, current_therapy);
    }
    else
    {
        current_therapy.push_back(NULL);

        for (Antivirals::const_iterator it = antivirals.begin(); it != antivirals.end(); ++it)
        {
            if (it->apply(initial_alignment, new_alignment))
            {
                current_therapy.back() = &(*it);
                generate_therapy(new_alignment, antivirals, current_therapy);
            }
        }

        current_therapy.pop_back();
    }
}

void generate_therapies(const SparseAlignment& initial_alignment, const Antivirals& antivirals)
{
    Therapy therapy;
    generate_therapy(initial_alignment, antivirals, therapy);
}

bool load_antivirals(const char* filename, Antivirals& antivirals)
{
    //bool ret;
    //ParseErrorData ped;

    FXParserLOVMethod<Antiviral, Antiviral&, Parameters2<size_t, const string&> > resistanceParser("resistance", &Antiviral::add_resistance);
    resistanceParser.parameters.set_param0("pos");
    resistanceParser.parameters.set_param1("aminos");

    FXParserNonDef<Antiviral, Constructor1<Antiviral, string> > antiviralParser("antiviral");
    antiviralParser.constructor.set_param0("id");
    add_lov_child(antiviralParser, resistanceParser);

    FXPRootParserNonDef<Antivirals> antiviralsParser("antivirals", antiviralParser);

    return parse_xml(filename, antivirals, antiviralsParser, /*ped,*/ Warn, Warn);
}

int main(int argc, char* argv[])
{
    SparseAlignment alignment;
    Antivirals antivirals;

    if (load_antivirals("antivirals.xml", antivirals))
    {
        generate_therapies(alignment, antivirals);
        return 0;
    }
    else
        return 1;
}


// -----------------------------
bool Antiviral::apply(const SparseAlignment& alig_original, SparseAlignment& alig_resulting) const
{
    bool applicable(false);
    SparseSequencesList resistant_sequences; // partial expanded alignment for each sequence

    //alig_resulting.clear();

    // for each sequence:
    for (SparseAlignment::const_iterator seq = alig_original.begin(); seq != alig_original.end(); ++seq)
    {
        if (apply(*seq, resistant_sequences))
        {
            // do alig_resulting = resistant_sequences U partial_alignment
            // set applicable = true if there is some new sequence (that is, partial_alignment - original) != EMPTY.
            for (SparseSequencesList::const_iterator new_seq = resistant_sequences.begin(); new_seq != resistant_sequences.end(); ++new_seq)
            {
                if (!contains(alig_original, *new_seq))
                {
                    // if first time...
                    if (!applicable)
                        alig_resulting = alig_original;

                    alig_resulting.insert(*new_seq);
                    applicable = true;
                }
            }
        }
        //else
        //alig_resulting.insert(*seq);    // this sequence is not affected by the antiviral.
    }

    return applicable;
}

bool Antiviral::apply(const SparseSequence& sequence, SparseSequencesList& resistant_sequences) const
{
    // If it turns out that is applicable, then generate a new sequence with each non-existant resistant position.
    const bool applicable = applicable_to(sequence);

    resistant_sequences.clear();

    if (applicable)
    {
        // TODO: There should create a COMBINATORY of all the possible induced resistant sequences. However,
        // this version only takes 1 mutation per sequence.
        for (list<ResistancePosition>::const_iterator it = resistances.begin(); it != resistances.end(); ++it)
        {
            if (it->applicable(sequence))
            {
                for (string::const_iterator amino = it->aminoacids.begin(); amino != it->aminoacids.end(); ++amino)
                    resistant_sequences.push_back(SparseSequence(sequence).set_amino(it->position, *amino));
            }
        }
    }

    return applicable;
}

bool Antiviral::applicable_to(const SparseSequence& sequence) const
{
    bool applicable;
    list<ResistancePosition>::const_iterator it;

    // The antiviral can be applicable if NONE of the resistance positions exist in the sequence.
    // IOW, if at least ONE resistance position exists, the antiviral is not applicable.

    it = resistances.begin();
    do
    {
        applicable = it->applicable(sequence);
        ++it;
    }
    while (applicable && it != resistances.end());

    return applicable;
}

